package com.aimei.wepeaker.audio

import android.content.Context
import android.media.MediaExtractor
import android.media.MediaFormat
import android.util.Log
import org.json.JSONArray
import org.json.JSONObject
import java.io.Closeable
import java.io.IOException
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.nio.FloatBuffer
import java.nio.ShortBuffer
import kotlin.math.*
import ai.onnxruntime.*
import org.apache.commons.math3.complex.Complex
import org.apache.commons.math3.transform.DftNormalization
import org.apache.commons.math3.transform.FastFourierTransformer
import org.apache.commons.math3.transform.TransformType

/**
 * Audio purity detector for Android based on the Python implementation.
 * Uses ONNX Runtime for model inference.
 */
class SpeakerPurityDetector(private val context: Context) : Closeable {
    
    private val TAG = "SpeakerPurityDetector"
    private var session: OrtSession? = null
    private var melBasis: Array<FloatArray>? = null
    private var window: FloatArray? = null
    private val resampleRate = 16000
    
    companion object {
        private const val ONNX_MODEL_PATH = "cnceleb_resnet34.onnx"
        private const val MEL_BASIS_PATH = "wepeaker_mel_basis.json"
        private const val WINDOW_PATH = "wepeaker_window.json"
        
        // Load the native ONNX Runtime library
        init {
            try {
                System.loadLibrary("onnxruntime")
            } catch (e: UnsatisfiedLinkError) {
                Log.e("SpeakerPurityDetector", "Failed to load onnxruntime native library: ${e.message}")
            }
        }
    }
    
    init {
        Log.d(TAG, "Initializing SpeakerPurityDetector")
        try {
            initializeModel()
            loadMelBasisAndWindow()
        } catch (e: Exception) {
            Log.e(TAG, "Error during initialization: ${e.message}")
            throw IOException("Failed to initialize: ${e.message}")
        }
    }
    
    private fun initializeModel() {
        // Create session options
        try {
            val sessionOptions = OrtSession.SessionOptions().apply {
                setInterOpNumThreads(1)
                setIntraOpNumThreads(1)
                setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT)
            }
            
            // Load model from assets
            context.assets.open(ONNX_MODEL_PATH).use { inputStream ->
                val modelBytes = inputStream.readBytes()
                val env = OrtEnvironment.getEnvironment()
                session = env.createSession(modelBytes, sessionOptions)
                Log.d(TAG, "Model loaded successfully")
            }
        } catch (e: Exception) {
            Log.e(TAG, "Error initializing ONNX model", e)
            throw IOException("Failed to initialize ONNX model: ${e.message}")
        }
    }
    
    private fun loadMelBasisAndWindow() {
        // Load mel basis from JSON
        context.assets.open(MEL_BASIS_PATH).use { inputStream ->
            val jsonString = inputStream.bufferedReader().use { it.readText() }
            val jsonArray = JSONArray(jsonString)
            melBasis = Array(jsonArray.length()) { i ->
                val row = jsonArray.getJSONArray(i)
                FloatArray(row.length()) { j -> row.getDouble(j).toFloat() }
            }
            Log.d(TAG, "Mel basis loaded: ${melBasis?.size}x${melBasis?.getOrNull(0)?.size}")
        }
        
        // Load window from JSON
        context.assets.open(WINDOW_PATH).use { inputStream ->
            val jsonString = inputStream.bufferedReader().use { it.readText() }
            val jsonArray = JSONArray(jsonString)
            window = FloatArray(jsonArray.length()) { i -> jsonArray.getDouble(i).toFloat() }
            Log.d(TAG, "Window loaded: ${window?.size}")
        }
    }
    
    /**
     * Calculate audio purity from an asset file
     */
    fun detectPurity(audioFileName: String): Float {
        Log.d(TAG, "Detecting purity for file: $audioFileName")
        
        // Load audio data from asset
        val audioData = loadAudioFromAsset(audioFileName)
        
        // Calculate purity
        return runPurityDetection(audioData)
    }
    
    /**
     * Load audio from asset and return 16kHz mono PCM samples
     */
    private fun loadAudioFromAsset(audioFileName: String): FloatArray {
        Log.d(TAG, "Loading audio from asset: $audioFileName")
        
        context.assets.openFd(audioFileName).use { afd ->
            val extractor = MediaExtractor()
            extractor.setDataSource(afd.fileDescriptor, afd.startOffset, afd.length)
            
            var sampleRate = 0
            var channelCount = 0
            
            // Select audio track
            for (i in 0 until extractor.trackCount) {
                val format = extractor.getTrackFormat(i)
                val mime = format.getString(MediaFormat.KEY_MIME)
                if (mime?.startsWith("audio/") == true) {
                    extractor.selectTrack(i)
                    sampleRate = format.getInteger(MediaFormat.KEY_SAMPLE_RATE)
                    channelCount = format.getInteger(MediaFormat.KEY_CHANNEL_COUNT)
                    break
                }
            }
            
            if (sampleRate == 0) {
                throw IOException("No audio track found in the file")
            }
            
            Log.d(TAG, "Audio info: rate=$sampleRate, channels=$channelCount")
            
            // Read all samples
            val samples = mutableListOf<Float>()
            val maxBuffer = ByteBuffer.allocateDirect(1024 * 2) // 1024 samples
            maxBuffer.order(ByteOrder.nativeOrder())
            
            while (true) {
                maxBuffer.clear()
                val sampleSize = extractor.readSampleData(maxBuffer, 0)
                if (sampleSize <= 0) break
                
                maxBuffer.rewind()
                val shortBuffer = maxBuffer.asShortBuffer()
                val tempArray = ShortArray(sampleSize / 2)
                shortBuffer.get(tempArray)
                
                for (j in tempArray.indices) {
                    samples.add(tempArray[j] / 32768.0f)
                }
                
                extractor.advance()
            }
            
            extractor.release()
            
            // Convert to mono if stereo
            val monoSamples = if (channelCount > 1) {
                val mono = FloatArray(samples.size / channelCount)
                for (i in mono.indices) {
                    var sum = 0.0f
                    for (ch in 0 until channelCount) {
                        sum += samples[i * channelCount + ch]
                    }
                    mono[i] = sum / channelCount
                }
                mono
            } else {
                samples.toFloatArray()
            }
            
            // Resample to 16kHz if needed
            return if (sampleRate != resampleRate) {
                simpleResample(monoSamples, sampleRate, resampleRate)
            } else {
                monoSamples
            }
        }
    }
    
    /**
     * Simple resampling method
     */
    private fun simpleResample(audio: FloatArray, originalRate: Int, targetRate: Int): FloatArray {
        val ratio = targetRate.toFloat() / originalRate
        val newLength = (audio.size * ratio).toInt()
        val resampled = FloatArray(newLength)
        
        for (i in resampled.indices) {
            val srcIdx = (i / ratio).toInt()
            val nextIdx = min(srcIdx + 1, audio.size - 1)
            val fraction = (i / ratio) - srcIdx
            
            resampled[i] = audio[srcIdx] * (1 - fraction) + audio[nextIdx] * fraction
        }
        
        return resampled
    }
    
    /**
     * Extract embedding from audio using ONNX model
     */
    private fun extractEmbeddingFromPcm(wav: FloatArray): FloatArray {
        Log.d(TAG, "Input audio check - shape: ${wav.size}, sum: ${wav.sum()}")
        
        val embeddings = FloatArray(256)
        var i = 0
        while (i < wav.size) {
            val wavPiece = wav.sliceArray(i until min(i + 16000 * 3, wav.size))
            val weight = wavPiece.size / 48000f
            
            // Pad if necessary
            val paddedWav = if (wavPiece.size < 16000 * 3) {
                FloatArray(16000 * 3).also { wavPiece.copyInto(it) }
            } else {
                wavPiece
            }
            
            // Scale to int16 range
            val scaledWav = paddedWav.map { it * 32768f }.toFloatArray()
            Log.d(TAG, "Pre-fbank audio check - shape: ${scaledWav.size}, sum: ${scaledWav.sum()}")
            
            // Calculate fbank features
            val feats = fbankNumpy(scaledWav)
            Log.d(TAG, "Fbank results check - shape: ${feats.size}x${feats[0].size}, sum: ${feats.sumByDouble { it.sum().toDouble() }}")
            
            // Mean normalization
            val means = FloatArray(feats[0].size)
            for (j in feats.indices) {
                for (k in feats[j].indices) {
                    means[k] += feats[j][k]
                }
            }
            for (j in means.indices) {
                means[j] /= feats.size
            }
            
            for (j in feats.indices) {
                for (k in feats[j].indices) {
                    feats[j][k] -= means[k]
                }
            }
            
            // Prepare input tensor
            val inputTensor = OnnxTensor.createTensor(
                OrtEnvironment.getEnvironment(),
                FloatBuffer.wrap(feats.flatMap { it.toList() }.toFloatArray()),
                longArrayOf(1, feats.size.toLong(), feats[0].size.toLong())
            )
            
            // Run inference
            val inputs = mapOf("feats" to inputTensor)
            session?.run(inputs)?.use { results ->
                val outputTensor = results["embs"] as OnnxTensor
                val outputData = outputTensor.value as Array<*>
                val embedding = (outputData[0] as Array<*>).map { (it as Number).toFloat() }.toFloatArray()
                
                for (j in embedding.indices) {
                    embeddings[j] += embedding[j] * weight
                }
                
                Log.d(TAG, "Embedding check - shape: ${embedding.size}, sum: ${embedding.sum()}")
            }
            
            i += 16000 * 3
        }
        
        return embeddings
    }
    
    /**
     * Fbank feature extraction (similar to Python's librosa.feature.melspectrogram)
     */
    private fun fbankNumpy(waveform: FloatArray): Array<FloatArray> {
        val melBasisLocal = melBasis ?: throw IllegalStateException("Mel basis not initialized")
        val windowLocal = window ?: throw IllegalStateException("Window not initialized")
        
        val preemphasisCoef = 0.97f
        val windowShift = 160
        val windowSize = 400
        val paddedWindowSize = 512
        
        // Step 1: Create strided input
        val numSamples = waveform.size
        val m = 1 + (numSamples - windowSize) / windowShift
        val stridedInput = Array(m) { i ->
            FloatArray(windowSize) { j ->
                waveform[i * windowShift + j]
            }
        }
        Log.d(TAG, "Fbank step 1 check - shape: ${stridedInput.size}x${stridedInput[0].size}, sum: ${stridedInput.sumByDouble { it.sum().toDouble() }}")
        
        // Step 2: Mean subtraction
        for (i in stridedInput.indices) {
            val mean = stridedInput[i].sum() / stridedInput[i].size
            for (j in stridedInput[i].indices) {
                stridedInput[i][j] -= mean
            }
        }
        Log.d(TAG, "Fbank step 2 check - shape: ${stridedInput.size}x${stridedInput[0].size}, sum: ${stridedInput.sumByDouble { it.sum().toDouble() }}")
        
        // Step 3 & 4: Preemphasis
        val stridedInputPreemphasis = Array(stridedInput.size) { i ->
            FloatArray(stridedInput[i].size) { j ->
                if (j == 0) {
                    stridedInput[i][j]
                } else {
                    stridedInput[i][j] - preemphasisCoef * stridedInput[i][j - 1]
                }
            }
        }
        Log.d(TAG, "Fbank step 4 check - shape: ${stridedInputPreemphasis.size}x${stridedInputPreemphasis[0].size}, sum: ${stridedInputPreemphasis.sumByDouble { it.sum().toDouble() }}")
        
        // Step 5: Apply window
        val stridedInputWindowed = Array(stridedInputPreemphasis.size) { i ->
            FloatArray(stridedInputPreemphasis[i].size) { j ->
                stridedInputPreemphasis[i][j] * windowLocal[j]
            }
        }
        Log.d(TAG, "Fbank step 5 check - shape: ${stridedInputWindowed.size}x${stridedInputWindowed[0].size}, sum: ${stridedInputWindowed.sumByDouble { it.sum().toDouble() }}")
        
        // Step 6: Padding
        val paddingRight = paddedWindowSize - windowSize
        val stridedInputPadded = Array(stridedInputWindowed.size) { i ->
            FloatArray(paddedWindowSize) { j ->
                if (j < windowSize) stridedInputWindowed[i][j] else 0f
            }
        }
        Log.d(TAG, "Fbank step 6 check - shape: ${stridedInputPadded.size}x${stridedInputPadded[0].size}, sum: ${stridedInputPadded.sumByDouble { it.sum().toDouble() }}")
        
        // Step 7: FFT
        val fft = FastFourierTransformer(DftNormalization.STANDARD)
        val outputSize = paddedWindowSize / 2 + 1
        val spectrum = Array(stridedInputPadded.size) { i ->
            val complexInput = stridedInputPadded[i].map { Complex(it.toDouble(), 0.0) }.toTypedArray()
            val fftResult = fft.transform(complexInput, TransformType.FORWARD)
            val magnitudes = FloatArray(outputSize) { j ->
                val magnitude = sqrt(fftResult[j].real * fftResult[j].real + fftResult[j].imaginary * fftResult[j].imaginary).toFloat()
                magnitude * magnitude // Power spectrum
            }
            magnitudes
        }
        Log.d(TAG, "Fbank step 7 check - shape: ${spectrum.size}x${spectrum[0].size}, sum: ${spectrum.sumByDouble { it.sum().toDouble() }}")
        
        // Step 8: Apply mel filterbank
        val melEnergies = Array(spectrum.size) { i ->
            FloatArray(melBasisLocal.size) { j ->
                var sum = 0f
                for (k in spectrum[i].indices) {
                    sum += spectrum[i][k] * melBasisLocal[j][k]
                }
                max(sum, 1.1921e-7f)
            }
        }
        
        // Apply log
        for (i in melEnergies.indices) {
            for (j in melEnergies[i].indices) {
                melEnergies[i][j] = ln(melEnergies[i][j])
            }
        }
        
        Log.d(TAG, "Fbank step 8 (final) check - shape: ${melEnergies.size}x${melEnergies[0].size}, sum: ${melEnergies.sumByDouble { it.sum().toDouble() }}")
        return melEnergies
    }
    
    /**
     * Calculate RMS (Root Mean Square) values
     */
    private fun toRms(wav: FloatArray, windowSize: Int = 1024, hopSize: Int = 160): FloatArray {
        val fNums = wav.size / hopSize
        val result = FloatArray(fNums)
        
        for (i in 0 until fNums) {
            var sum = 0f
            var count = 0
            
            for (j in max(0, i * hopSize - windowSize / 2) until min(wav.size, i * hopSize + windowSize / 2)) {
                sum += wav[j] * wav[j]
                count++
            }
            
            result[i] = sqrt(sum / count)
        }
        
        return result
    }
    
    /**
     * Calculate cosine similarity
     */
    private fun cosineSimilarity(a: FloatArray, b: FloatArray): Float {
        var dotProduct = 0f
        var normA = 0f
        var normB = 0f
        
        for (i in a.indices) {
            dotProduct += a[i] * b[i]
            normA += a[i] * a[i]
            normB += b[i] * b[i]
        }
        
        return dotProduct / (sqrt(normA) * sqrt(normB))
    }
    
    /**
     * Run the purity detection algorithm
     */
    private fun runPurityDetection(wav: FloatArray): Float {
        val rmss = toRms(wav)
        val setenceEmbeddings = mutableListOf<FloatArray>()
        val setenceRmsWeight = mutableListOf<Float>()
        
        // Process audio in chunks
        var i = 0
        while (i < wav.size) {
            val wavPiece = wav.sliceArray(i until min(i + 16000 * 6, wav.size))
            if (wavPiece.size < 1000) break // Skip too small pieces
            
            val embedding = extractEmbeddingFromPcm(wavPiece)
            
            // Normalize embedding
            val norm = sqrt(embedding.fold(0f) { acc, f -> acc + f * f })
            val normalizedEmbedding = embedding.map { it / norm }.toFloatArray()
            
            setenceEmbeddings.add(normalizedEmbedding)
            
            // Calculate RMS weight
            val startFrame = i / 160
            val endFrame = min(startFrame + 600, rmss.size - 1)
            if (startFrame < endFrame) {
                val rmsWeight = rmss.sliceArray(startFrame until endFrame).average().toFloat()
                setenceRmsWeight.add(rmsWeight)
            } else {
                setenceRmsWeight.add(0.001f) // Avoid division by zero
            }
            
            i += 16000 * 6
        }
        
        if (setenceEmbeddings.isEmpty()) {
            Log.w(TAG, "No embeddings were extracted")
            return 0.5f
        }
        
        // Calculate similarity matrix
        val simMatrix = Array(setenceEmbeddings.size) { i ->
            FloatArray(setenceEmbeddings.size) { j ->
                cosineSimilarity(setenceEmbeddings[i], setenceEmbeddings[j])
            }
        }
        
        // Find max cluster
        var maxClusterIdx = 0
        var maxClusterSize = 0
        for (i in simMatrix.indices) {
            val clusterSize = simMatrix[i].count { it > 0.6f }
            if (clusterSize > maxClusterSize) {
                maxClusterSize = clusterSize
                maxClusterIdx = i
            }
        }
        
        // Calculate kernel embedding
        var kernelEmbedding = FloatArray(setenceEmbeddings[0].size)
        val initialCluster = mutableListOf<Int>()
        
        for (i in simMatrix.indices) {
            if (simMatrix[maxClusterIdx][i] > 0.6f) {
                initialCluster.add(i)
            }
        }
        
        for (idx in initialCluster) {
            for (i in setenceEmbeddings[idx].indices) {
                kernelEmbedding[i] += setenceEmbeddings[idx][i] / initialCluster.size
            }
        }
        
        // Normalize kernel embedding
        val kernelNorm = sqrt(kernelEmbedding.fold(0f) { acc, f -> acc + f * f })
        for (i in kernelEmbedding.indices) {
            kernelEmbedding[i] /= kernelNorm
        }
        
        // Refine kernel embedding (5 iterations)
        for (iter in 0 until 5) {
            val kernelDis = setenceEmbeddings.map { cosineSimilarity(kernelEmbedding, it) }
            val clusterMembers = kernelDis.mapIndexedNotNull { idx, sim -> if (sim > 0.6f) idx else null }
            
            Log.d(TAG, "Iteration $iter, cluster size: ${clusterMembers.size}/${setenceEmbeddings.size}")
            
            if (clusterMembers.isEmpty()) break
            
            // Reset kernel embedding
            kernelEmbedding = FloatArray(setenceEmbeddings[0].size)
            
            // Calculate new kernel embedding
            for (idx in clusterMembers) {
                for (i in setenceEmbeddings[idx].indices) {
                    kernelEmbedding[i] += setenceEmbeddings[idx][i] / clusterMembers.size
                }
            }
            
            // Normalize kernel embedding
            val newNorm = sqrt(kernelEmbedding.fold(0f) { acc, f -> acc + f * f })
            for (i in kernelEmbedding.indices) {
                kernelEmbedding[i] /= newNorm
            }
        }
        
        // Final similarity calculation
        val kernelDis = setenceEmbeddings.map { cosineSimilarity(kernelEmbedding, it) }
        
        // Calculate purity
        var numerator = 0f
        var denominator = 0f
        
        for (i in kernelDis.indices) {
            if (kernelDis[i] > 0.6f) {
                numerator += setenceRmsWeight[i]
            }
            denominator += setenceRmsWeight[i]
        }
        
        val purity = if (denominator > 0f) numerator / denominator else 0.5f
        Log.d(TAG, "Final purity: $purity")
        
        return purity
    }
    
    override fun close() {
        Log.d(TAG, "Closing SpeakerPurityDetector")
        session?.close()
    }
}